import os
import pickle
from flask import Flask, request, jsonify, render_template, redirect, url_for
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import pandas as pd


# 配置
MODEL_PATH = 'model/model.pkl'
DATA_PATH = 'data/diabetes.csv'
# 确保目录存在
os.makedirs('model', exist_ok=True)
os.makedirs('data', exist_ok=True)

app = Flask(__name__)
model = None

# 特征处理函数（针对糖尿病数据集）
def preprocess(data, is_training=False):
    df = pd.DataFrame(data)
    
    # 1. 将 0 视为缺失值的列填充为均值
    zero_as_nan_cols = ['plas', 'pres', 'skin', 'insu', 'mass']
    for col in zero_as_nan_cols:
        if col in df.columns:
            df[col] = df[col].replace(0, pd.NA)
            mean_val = df[col].mean(skipna=True)
            df[col] = df[col].fillna(mean_val)
    
    # 2. 如果是训练数据，处理目标列
    if is_training and 'class' in df.columns:
        df['class'] = df['class'].map({'tested_negative': 0, 'tested_positive': 1})
    
    return df

# 加载模型
def load_model():
    global model
    if os.path.exists(MODEL_PATH):
        with open(MODEL_PATH, 'rb') as f:
            model = pickle.load(f)
        return True
    else:
        model = None
        return False

# 保存模型
def save_model(m):
    with open(MODEL_PATH, 'wb') as f:
        pickle.dump(m, f)

# API: 预测
@app.route('/predict', methods=['POST'])
def predict():
    global model
    if model is None:
        return jsonify({'error': '没有加载模型，请先训练模型'}), 400
    
    try:
        data = request.get_json()
        # 验证输入数据
        required_features = ['preg', 'plas', 'pres', 'skin', 'insu', 'mass', 'pedi', 'age']
        for item in data:
            for feature in required_features:
                if feature not in item:
                    return jsonify({'error': f'缺少必要特征: {feature}'}), 400
        
        X = preprocess(data)
        pred = model.predict(X)
        # 将预测结果转换为有意义的文本
        pred_labels = ['tested_negative' if p == 0 else 'tested_positive' for p in pred]
        return jsonify({'prediction': pred_labels})
    except Exception as e:
        return jsonify({'error': str(e)}), 500


# API: 训练
@app.route('/train', methods=['GET'])
def train():
    if not os.path.exists(DATA_PATH):
        return jsonify({'error': f'未找到数据集，请确保{DATA_PATH}存在'}), 400
    
    try:
        df = pd.read_csv(DATA_PATH)
        # 检查必要的列是否存在
        required_columns = ['preg', 'plas', 'pres', 'skin', 'insu', 'mass', 'pedi', 'age', 'class']
        for col in required_columns:
            if col not in df.columns:
                return jsonify({'error': f'数据集中缺少必要的列: {col}'}), 400
        
        # 预处理数据
        X = preprocess(df.drop('class', axis=1))
        y = df['class'].map({'tested_negative': 0, 'tested_positive': 1})
        
        # 分割训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        
        # 训练模型
        clf = RandomForestClassifier(n_estimators=100, random_state=42)
        clf.fit(X_train, y_train)
        
        # 在测试集上评估
        y_pred = clf.predict(X_test)
        acc = accuracy_score(y_test, y_pred)
        
        # 保存并加载模型
        save_model(clf)
        load_model()
        
        return jsonify({
            'result': '模型训练完成', 
            'test_accuracy': float(acc),
            'message': f'模型准确率: {acc:.2f}'
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

# API: 清除模型
@app.route('/wipe', methods=['GET'])
def wipe():
    global model
    if os.path.exists(MODEL_PATH):
        os.remove(MODEL_PATH)
        model = None
        return jsonify({'result': '模型已清除'})
    else:
        return jsonify({'result': '没有模型可清除'}), 400

# Web页面
@app.route('/web', methods=['GET', 'POST'])
def web():
    global model
    result = None
    error = None
    
    if request.method == 'POST':
        try:
            # 获取表单数据
            data = {
                'preg': float(request.form.get('preg', 0)),
                'plas': float(request.form.get('plas', 0)),
                'pres': float(request.form.get('pres', 0)),
                'skin': float(request.form.get('skin', 0)),
                'insu': float(request.form.get('insu', 0)),
                'mass': float(request.form.get('mass', 0)),
                'pedi': float(request.form.get('pedi', 0)),
                'age': float(request.form.get('age', 0))
            }
            
            # 验证输入
            for key, value in data.items():
                if value < 0:
                    raise ValueError(f'{key}不能为负数')
            
            if model:
                X = preprocess([data])
                pred = model.predict(X)[0]
                result = f'预测结果: {"患有糖尿病" if pred == 1 else "未患有糖尿病"}'
            else:
                error = '模型未加载，请先训练模型'
        except ValueError as ve:
            error = f'输入错误: {str(ve)}'
        except Exception as e:
            error = f'处理错误: {str(e)}'
    
    return render_template('web.html', result=result, error=error)

# 首页跳转
@app.route('/')
def index():
    return redirect(url_for('web'))

if __name__ == '__main__':
    import sys
    port = 5000
    if len(sys.argv) > 1:
        try:
            port = int(sys.argv[1])
        except:
            pass
    # 尝试加载已保存的模型
    load_model()
    app.run(host='0.0.0.0', port=port, debug=False)  # 生产环境关闭debug
